import os
import numpy as np

from model import TGLCGAN
from utils import pp

import tensorflow as tf

flags = tf.app.flags
flags.DEFINE_integer("epoch", 25, "Epoch to train [25]")
flags.DEFINE_float("learning_rate", 0.0002, "Learning rate of for adam [0.0002]")
flags.DEFINE_integer("sample_height", 32, "The height of sample data")
flags.DEFINE_integer("sample_width", 32, "The width of sample data")
flags.DEFINE_integer("sample_depth", 128, "The depth of sample data")
flags.DEFINE_float("beta1", 0.5, "Momentum term of adam [0.5]")
flags.DEFINE_integer("train_size", np.inf, "The size of train images [np.inf]")
flags.DEFINE_integer("batch_size", 16, "The size of batch images [64]")
flags.DEFINE_string("dataset", "seis_1", "The name of dataset [sies_1, seis_2]")
flags.DEFINE_string("checkpoint_dir", "checkpoint", "Directory name to save the checkpoints [checkpoint]")
flags.DEFINE_boolean("train", True, "True for training, False for testing [False]")
FLAGS = flags.FLAGS

def main(_):
    pp.pprint(flags.FLAGS.__flags)

    if not os.path.exists(FLAGS.checkpoint_dir):
        os.makedirs(FLAGS.checkpoint_dir)

    # gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
    run_config = tf.ConfigProto()
    run_config.gpu_options.allow_growth = True

    with tf.Session(config=run_config) as sess:
        gan = TGLCGAN(
            sess,
            batch_size=FLAGS.batch_size,
            dataset_name=FLAGS.dataset,
            sample_shape=[FLAGS.sample_height, FLAGS.sample_width, FLAGS.sample_depth],
            checkpoint_dir=FLAGS.checkpoint_dir)

        if FLAGS.train:
            gan.train(FLAGS)
        else:
            gan.compeletion(FLAGS)

if __name__ == '__main__':
    tf.app.run()
